use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use actix_web::{dev::Service, dev::ServiceRequest, dev::ServiceResponse, Error};
use actix_web::dev::{ServiceFactory, Transform};

pub struct KeepAliveMiddleware<S> {
    service: S,
}

impl<S, B> Service<ServiceRequest> for KeepAliveMiddleware<S>
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
    S::Future: 'static,
{
    type Response = ServiceResponse<B>;
    type Error = Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>>>>;

    fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.service.poll_ready(cx)
    }

    fn call(&self, req: ServiceRequest) -> Self::Future {
        let fut = self.service.call(req);
        Box::pin(async move {
            let res = fut.await?;
            // 在这里可以添加保活逻辑
            Ok(res)
        })
    }
}

pub struct KeepAlive;

impl<S, B> Transform<S, ServiceRequest> for KeepAlive
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
    S::Future: 'static,
{
    type Response = ServiceResponse<B>;
    type Error = Error;
    type Transform = KeepAliveMiddleware<S>;
    type InitError = ();
    type Future = Pin<Box<dyn Future<Output = Result<Self::Transform, Self::InitError>>>>;

    fn new_transform(&self, service: S) -> Self::Future {
        Box::pin(async move { Ok(KeepAliveMiddleware { service }) })
    }
}

pub fn keep_alive() -> KeepAlive {
    KeepAlive
}